# -*- coding: utf-8 -*-
# @Time    : 2025/10/5 13:53
# @Author  : 楚江涛
# @FileName: state_memory.py
# @Software: PyCharm


"""
apps/workflow/infrastructure/state_memory.py
----------------------------------
基于内存的执行状态存储（StateStore 实现，完整版）。

用途：
- 维护 /executions/{exec_id}/status 所需的“可查询快照”；
- 记录执行级状态（RUNNING/SUCCEEDED/FAILED/CANCELLED）、节点运行/完成列表、错误信息与进度；
- 轻量、无外部依赖，适合本地开发或单进程部署。

设计要点：
- 线程/协程安全：使用 threading.RLock 保证多线程/异步上下文下的数据一致性；
- 幂等：所有变更操作具备幂等语义；
- 快照友好：snapshot() 返回浅拷贝，避免外部修改内部状态；
- 可扩展字段：记录 created_at / updated_at / finished_at，用于观测；
- 进度：set_progress() 可选，未设置则不返回该字段；
- 错误：errors 映射 node_id -> 错误文本（可配置最大长度，防止爆内存）。

注意：
- 多进程/多实例场景请替换为 Redis/DB 实现。
"""

from __future__ import annotations

import threading
import time
from typing import Any, Dict

from domain.enums import ExecStatus
from execution.interfaces import StateStore


class MemoryStateStore(StateStore):
    """
    内存状态存储实现。

    内部结构示例（_data[exec_id]）：
    {
        "status": "RUNNING",
        "running": ["nodeA"],
        "finished": ["start_0"],
        "errors": {"nodeX": "错误信息"},
        "progress": 0.42,                 # 可选
        "created_at": 1710000000.0,       # unix ts
        "updated_at": 1710000123.4,
        "finished_at": 1710000456.7,      # SUCCEEDED/FAILED/CANCELLED 时写入
    }
    """

    # 错误文本的最大保留长度（过长会截断）
    MAX_ERROR_TEXT = 2000

    def __init__(self) -> None:
        # 执行 ID -> 状态字典
        self._data: Dict[str, Dict[str, Any]] = {}
        # 读写锁
        self._lock = threading.RLock()

    # ============================================================
    # 执行级别
    # ============================================================

    def init(self, exec_id: str) -> None:
        """
        初始化执行快照（PENDING）。
        幂等：若已存在则重置为 PENDING 并清空集合。
        """
        now = time.time()
        with self._lock:
            self._data[exec_id] = {
                "status": ExecStatus.PENDING.value,
                "running": [],
                "finished": [],
                "errors": {},
                "created_at": now,
                "updated_at": now,
                # "progress": 0.0,   # 可选：不默认返回，避免误导
                # "finished_at": None
            }

    def update_status(self, exec_id: str, status: ExecStatus) -> None:
        """
        更新执行状态（RUNNING/SUCCEEDED/FAILED/CANCELLED）。
        结束态（SUCCEEDED/FAILED/CANCELLED）会写入 finished_at。
        """
        with self._lock:
            state = self._require(exec_id)
            state["status"] = status.value
            state["updated_at"] = time.time()
            if status in (ExecStatus.SUCCEEDED, ExecStatus.FAILED, ExecStatus.CANCELLED):
                state["finished_at"] = state.get("finished_at") or state["updated_at"]

    def set_progress(self, exec_id: str, progress: float) -> None:
        """
        设置进度（0.0 ~ 1.0）。超出范围会裁剪。
        """
        p = 0.0 if progress < 0.0 else (1.0 if progress > 1.0 else float(progress))
        with self._lock:
            state = self._require(exec_id)
            state["progress"] = p
            state["updated_at"] = time.time()

    # ============================================================
    # 节点级别
    # ============================================================

    def mark_node_running(self, exec_id: str, node_id: str) -> None:
        """
        标记节点为运行中：加入 running（若已存在则忽略）。
        """
        with self._lock:
            st = self._require(exec_id)
            running = st["running"]
            if node_id not in running:
                running.append(node_id)
            # 若之前被标记为 finished（例如重试逻辑），可选择移除；此处保留已完成记录
            st["updated_at"] = time.time()

    def mark_node_finished(self, exec_id: str, node_id: str) -> None:
        """
        标记节点完成：从 running 移除，加入 finished（若已存在则忽略）。
        """
        with self._lock:
            st = self._require(exec_id)
            running = st["running"]
            finished = st["finished"]
            if node_id in running:
                running.remove(node_id)
            if node_id not in finished:
                finished.append(node_id)
            st["updated_at"] = time.time()

    def mark_node_failed(self, exec_id: str, node_id: str, error: str) -> None:
        """
        标记节点失败：完成标记 + 记录错误文本 + 执行状态置为 FAILED（幂等）。
        """
        err = (error or "").strip()
        if len(err) > self.MAX_ERROR_TEXT:
            err = err[: self.MAX_ERROR_TEXT] + " ...（已截断）"

        with self._lock:
            st = self._require(exec_id)
            # 完成标记
            running = st["running"]
            finished = st["finished"]
            if node_id in running:
                running.remove(node_id)
            if node_id not in finished:
                finished.append(node_id)
            # 错误记录
            st["errors"][node_id] = err
            # 执行状态 → FAILED（若此前为 CANCELLED，不强制覆盖）
            if st["status"] not in (ExecStatus.CANCELLED.value, ExecStatus.FAILED.value):
                st["status"] = ExecStatus.FAILED.value
            st["updated_at"] = time.time()

    # ============================================================
    # 读取与维护
    # ============================================================

    def snapshot(self, exec_id: str) -> Dict[str, Any]:
        """
        返回当前快照（浅拷贝），以防外部误改内部存储。
        未设置的可选字段（如 progress/finished_at）将被省略。
        """
        with self._lock:
            st = self._require(exec_id)
            snap = {
                "status": st["status"],
                "running": list(st["running"]),
                "finished": list(st["finished"]),
                "errors": dict(st["errors"]),
                "created_at": st["created_at"],
                "updated_at": st["updated_at"],
            }
            if "progress" in st:
                snap["progress"] = st["progress"]
            if "finished_at" in st:
                snap["finished_at"] = st["finished_at"]
            return snap

    # -----------------------------
    # 辅助/维护（可选对外公开）
    # -----------------------------

    def reset(self, exec_id: str) -> None:
        """
        将执行重置到 PENDING，清空运行/完成/错误/进度/完成时间。
        """
        with self._lock:
            st = self._require(exec_id)
            now = time.time()
            st["status"] = ExecStatus.PENDING.value
            st["running"].clear()
            st["finished"].clear()
            st["errors"].clear()
            st.pop("progress", None)
            st.pop("finished_at", None)
            st["updated_at"] = now

    def remove(self, exec_id: str) -> None:
        """
        删除一个执行的状态（不可恢复）。通常用于清理过期记录。
        """
        with self._lock:
            self._data.pop(exec_id, None)

    def exists(self, exec_id: str) -> bool:
        """判断某个执行是否存在状态记录"""
        with self._lock:
            return exec_id in self._data

    # ============================================================
    # 内部工具
    # ============================================================

    def _require(self, exec_id: str) -> Dict[str, Any]:
        """
        取出执行状态；不存在则抛出 KeyError（由上层决定如何处理）。
        """
        st = self._data.get(exec_id)
        if st is None:
            raise KeyError(f"State for exec_id={exec_id!r} not found.")
        return st

